"""
多视频流处理系统
从数据库读取视频流信息，处理多个视频流，并将结果存入Redis
"""

import cv2
import numpy as np
import yaml
import logging
import os
import time
import argparse
import threading
import base64
import json
import asyncio
from typing import Dict, List, Any, Optional, Tuple
from datetime import datetime

# 导入自定义模块
from detector import PersonDetector
from tracker import PersonTracker, TrackInfo
from region_manager import RegionManager, RegionConfig
from counter import PersonCounter
from video_processor import RTSPProcessor, FPSCounter
from db_manager import DBManager
from redis_manager import RedisManager
from webrtc_server import WebRTCServer

# 配置日志
logger = logging.getLogger(__name__)

class StreamProcessor:
    """单个视频流处理器"""
    
    def __init__(self, stream_info: Dict[str, Any], model_config: Dict[str, Any], 
                 tracking_config: Dict[str, Any], redis_manager: RedisManager):
        """
        初始化视频流处理器
        
        Args:
            stream_info: 视频流信息
            model_config: 模型配置
            tracking_config: 跟踪配置
            redis_manager: Redis管理器
        """
        self.stream_info = stream_info
        self.stream_id = str(stream_info['id'])
        self.camera_id = str(stream_info['camera_id'])
        self.rtsp_url = stream_info['rtsp_address']
        self.redis_key = f"het_img:{self.stream_id}:{self.camera_id}"
        self.redis_manager = redis_manager
        
        # 初始化区域
        self.region_points = stream_info['line_direction']
        
        # 初始化组件
        self.detector = None
        self.tracker = None
        self.region_manager = None
        self.counter = None
        self.video_processor = None
        
        # 运行状态
        self.is_running = False
        self.processing_thread = None
        self.frame_count = 0
        self.last_save_time = 0
        self.save_interval = 3.0  # 保存间隔（秒）
        
        # 用于WebRTC的最新处理后帧
        self.latest_processed_frame = None
        self.frame_lock = threading.Lock()
        
        # 初始化组件
        self._setup_components(model_config, tracking_config)
    
    def _setup_components(self, model_config: Dict[str, Any], tracking_config: Dict[str, Any]) -> bool:
        """
        设置处理组件
        
        Args:
            model_config: 模型配置
            tracking_config: 跟踪配置
            
        Returns:
            设置是否成功
        """
        try:
            # 初始化检测器
            self.detector = PersonDetector(
                model_path=model_config['model_path'],
                confidence_threshold=model_config['confidence_threshold'],
                iou_threshold=model_config['iou_threshold'],
                classes=model_config['classes']
            )
            
            # 初始化跟踪器
            self.tracker = PersonTracker(
                track_thresh=tracking_config['track_thresh'],
                track_buffer=tracking_config['track_buffer'],
                match_thresh=tracking_config['match_thresh'],
                frame_rate=15  # 固定帧率
            )
            
            # 初始化区域管理器
            region_config = RegionConfig(
                points=self.region_points,
                margin={'top': 50, 'bottom': 50, 'left': 50, 'right': 50},
                color=(0, 255, 0),
                thickness=2
            )
            self.region_manager = RegionManager(region_config)
            
            # 初始化计数器
            self.counter = PersonCounter(self.region_manager)
            
            # 初始化视频处理器
            self.video_processor = RTSPProcessor(
                rtsp_url=self.rtsp_url,
                fps_limit=15,
                resize_shape=(1280, 720)
            )
            
            logger.info(f"流 {self.stream_id} 组件设置完成")
            return True
            
        except Exception as e:
            logger.error(f"流 {self.stream_id} 设置组件失败: {e}")
            return False
    
    def start(self) -> bool:
        """
        启动视频流处理
        
        Returns:
            启动是否成功
        """
        if self.is_running:
            logger.warning(f"流 {self.stream_id} 已在运行中")
            return True
        
        # 启动视频捕获
        if not self.video_processor.start_capture():
            logger.error(f"流 {self.stream_id} 无法启动视频捕获")
            return False
        
        # 启动处理线程
        self.is_running = True
        self.processing_thread = threading.Thread(
            target=self._processing_loop,
            daemon=True
        )
        self.processing_thread.start()
        
        logger.info(f"流 {self.stream_id} 处理已启动")
        return True
    
    def stop(self):
        """停止视频流处理"""
        self.is_running = False
        
        # 等待处理线程结束
        if self.processing_thread and self.processing_thread.is_alive():
            self.processing_thread.join(timeout=2.0)
        
        # 停止视频捕获
        if self.video_processor:
            self.video_processor.stop_capture()
        
        logger.info(f"流 {self.stream_id} 处理已停止")
    
    def _processing_loop(self):
        """视频处理主循环"""
        logger.info(f"流 {self.stream_id} 处理循环已启动")
        
        fps_counter = FPSCounter()
        
        while self.is_running:
            try:
                # 获取视频帧
                frame = self.video_processor.get_frame(timeout=1.0)
                if frame is None:
                    continue
                
                # 处理帧
                processed_frame, count_status = self._process_frame(frame)
                
                # 更新帧计数
                self.frame_count += 1
                
                # 计算FPS
                fps = fps_counter.update()
                
                # 保存最新处理后的帧用于WebRTC
                with self.frame_lock:
                    self.latest_processed_frame = processed_frame.copy()
                
                # 检查是否需要保存到Redis
                current_time = time.time()
                if current_time - self.last_save_time >= self.save_interval:
                    # 保存到Redis
                    self._save_to_redis(processed_frame, count_status)
                    self.last_save_time = current_time
                    
                    logger.debug(f"流 {self.stream_id} 帧 {self.frame_count}: 人数={count_status['current_count']}, FPS={fps:.1f}")
                
            except Exception as e:
                logger.error(f"流 {self.stream_id} 处理帧时出错: {e}")
                time.sleep(1.0)  # 出错时等待一段时间
    
    def _process_frame(self, frame: np.ndarray) -> Tuple[np.ndarray, Dict[str, int]]:
        """
        处理单帧
        
        Args:
            frame: 输入帧
            
        Returns:
            处理后的帧和计数状态
        """
        # 检测人员
        detections = self.detector.detect(frame)
        
        # 跟踪人员
        tracked_detections = self.tracker.update(detections)
        track_infos = self.tracker.get_track_info(tracked_detections)
        
        # 更新计数器
        count_status = self.counter.update(track_infos)
        
        # 可视化
        annotated_frame = frame.copy()
        
        # 绘制区域
        annotated_frame = self.region_manager.draw_region(annotated_frame)
        
        # 绘制人员和计数信息
        self._draw_annotations(annotated_frame, tracked_detections, track_infos, count_status)
        
        return annotated_frame, count_status
    
    def _draw_annotations(self, frame: np.ndarray, tracked_detections, track_infos: List[TrackInfo], count_status: Dict[str, int]) -> np.ndarray:
        """
        在帧上绘制注释
        
        Args:
            frame: 输入帧
            tracked_detections: 跟踪检测结果
            track_infos: 跟踪信息列表
            count_status: 计数状态
            
        Returns:
            注释后的帧
        """
        # 创建一个与原始帧相同大小的透明覆盖层
        overlay = frame.copy()
        
        # 绘制每个人员的边界框和ID
        for i, track_info in enumerate(track_infos):
            # 获取边界框
            x1, y1, x2, y2 = track_info.bbox
            
            # 确定颜色（区域内为绿色，区域外为红色）
            color = (0, 255, 0) if self.region_manager.is_point_in_region(track_info.center) else (0, 0, 255)
            
            # 绘制边界框
            cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), color, 2)
            
            # 绘制ID
            cv2.putText(frame, f"ID:{track_info.track_id}", (int(x1), int(y1) - 10),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
            
            # 绘制轨迹
            if len(track_info.history) > 1:
                for j in range(1, len(track_info.history)):
                    cv2.line(frame, track_info.history[j-1], track_info.history[j], color, 2)
        
        # 绘制计数信息
        cv2.putText(frame, f"Enter: {count_status['enter_count']}", (10, 30),
                   cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
        cv2.putText(frame, f"Exit: {count_status['exit_count']}", (10, 60),
                   cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
        cv2.putText(frame, f"Current: {count_status['current_count']}", (10, 90),
                   cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
        
        # 添加流ID和时间戳
        current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        cv2.putText(frame, f"Stream ID: {self.stream_id}", (10, frame.shape[0] - 60),
                   cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
        cv2.putText(frame, f"Time: {current_time}", (10, frame.shape[0] - 30),
                   cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
        
        return frame
    
    def _save_to_redis(self, frame: np.ndarray, count_status: Dict[str, int]):
        """
        保存结果到Redis
        
        Args:
            frame: 处理后的帧
            count_status: 计数状态
        """
        try:
            # 将帧转换为Base64编码
            _, buffer = cv2.imencode('.jpg', frame)
            image_base64 = base64.b64encode(buffer).decode('utf-8')
            
            # 准备数据
            data = {
                'person_count': count_status['current_count'],  # 当前人数
                'helmet_count': 0,  # 安全帽数量，暂不支持
                'gap_person_and_helmet': count_status['current_count'],  # 未佩戴安全帽人数，暂不支持
                'image_base64': image_base64,
                'frame_count': self.frame_count,
                'stream_id': self.stream_id,
                'camera_id': self.camera_id
            }
            
            # 保存到Redis
            self.redis_manager.save_detection_result(self.redis_key, data)
            
        except Exception as e:
            logger.error(f"流 {self.stream_id} 保存到Redis失败: {e}")
    
    def get_status(self) -> Dict[str, Any]:
        """
        获取处理状态
        
        Returns:
            处理状态
        """
        video_stats = self.video_processor.get_stats() if self.video_processor else {}
        
        return {
            'stream_id': self.stream_id,
            'camera_id': self.camera_id,
            'rtsp_url': self.rtsp_url,
            'is_running': self.is_running,
            'frame_count': self.frame_count,
            'fps': video_stats.get('fps_actual', 0),
            'connected': video_stats.get('connected', False)
        }
        
    def get_latest_frame(self) -> Optional[np.ndarray]:
        """
        获取最新处理后的帧
        
        Returns:
            最新处理后的帧，如果没有则返回None
        """
        with self.frame_lock:
            if self.latest_processed_frame is not None:
                return self.latest_processed_frame.copy()
        return None


class MultiStreamProcessor:
    """多视频流处理系统"""
    
    def __init__(self, config_path: str = "config.yaml"):
        """
        初始化系统
        
        Args:
            config_path: 配置文件路径
        """
        self.config_path = config_path
        self.config = self._load_config()
        
        # 初始化日志
        self._setup_logging()
        
        # 数据库管理器
        self.db_manager = None
        
        # Redis管理器
        self.redis_manager = None
        
        # WebRTC服务器
        self.webrtc_server = None
        self.webrtc_runner = None
        self.webrtc_site = None
        
        # 流处理器
        self.stream_processors: Dict[str, StreamProcessor] = {}
        
        # 运行状态
        self.is_running = False
        self.monitor_thread = None
        
        logger.info("多视频流处理系统初始化完成")
    
    def _load_config(self) -> Dict[str, Any]:
        """加载配置文件"""
        try:
            with open(self.config_path, 'r', encoding='utf-8') as f:
                config = yaml.safe_load(f)
            return config
        except Exception as e:
            logger.error(f"加载配置文件失败: {e}")
            return self._get_default_config()
    
    def _get_default_config(self) -> Dict[str, Any]:
        """获取默认配置"""
        return {
            'database': {
                'host': 'localhost',
                'port': 3306,
                'user': 'root',
                'password': 'a361154',
                'database': 'xygk-factory-safe-v1'
            },
            'redis': {
                'host': 'localhost',
                'port': 6379,
                'password': 'lorytech',
                'db': 0
            },
            'webrtc': {
                'host': '0.0.0.0',
                'port': 8080,
                'enable': True
            },
            'model': {
                'model_path': 'yolo11n.pt',
                'confidence_threshold': 0.5,
                'iou_threshold': 0.45,
                'classes': [0]
            },
            'tracking': {
                'track_thresh': 0.25,
                'track_buffer': 30,
                'match_thresh': 0.8
            },
            'logging': {
                'level': 'INFO',
                'file_path': 'logs/system.log',
                'console': True
            },
            'system': {
                'monitor_interval': 60,  # 监控间隔（秒）
                'save_interval': 3.0     # 保存间隔（秒）
            }
        }
    
    def _setup_logging(self):
        """设置日志"""
        log_config = self.config.get('logging', {})
        level = getattr(logging, log_config.get('level', 'INFO'))
        
        # 创建日志目录
        log_file = log_config.get('file_path', 'logs/system.log')
        os.makedirs(os.path.dirname(log_file), exist_ok=True)
        
        # 配置日志
        logging.basicConfig(
            level=level,
            format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
            handlers=[
                logging.FileHandler(log_file, encoding='utf-8'),
                logging.StreamHandler() if log_config.get('console', True) else logging.NullHandler()
            ]
        )
    
    def setup(self) -> bool:
        """
        设置系统
        
        Returns:
            设置是否成功
        """
        try:
            # 初始化数据库管理器
            db_config = self.config['database']
            self.db_manager = DBManager(
                host=db_config['host'],
                port=db_config['port'],
                user=db_config['user'],
                password=db_config['password'],
                database=db_config['database']
            )
            
            # 初始化Redis管理器
            redis_config = self.config['redis']
            self.redis_manager = RedisManager(
                host=redis_config['host'],
                port=redis_config['port'],
                password=redis_config['password'],
                db=redis_config['db']
            )
            
            # 初始化WebRTC服务器
            webrtc_config = self.config['webrtc']
            if webrtc_config.get('enable', True):
                self.webrtc_server = WebRTCServer(
                    host=webrtc_config['host'],
                    port=webrtc_config['port']
                )
            
            logger.info("系统设置完成")
            return True
            
        except Exception as e:
            logger.error(f"设置系统失败: {e}")
            return False
    
    # 删除异步启动方法，改用同步方法
    
    def _webrtc_server_thread(self):
        """WebRTC服务器线程"""
        try:
            # 直接使用同步方法启动服务器
            if self.webrtc_server:
                logger.info(f"在线程中启动WebRTC服务器: {self.webrtc_server.host}:{self.webrtc_server.port}")
                self.webrtc_server.start()
        except Exception as e:
            logger.error(f"WebRTC服务器出错: {e}")
    
    def start(self) -> bool:
        """
        启动系统
        
        Returns:
            启动是否成功
        """
        if self.is_running:
            logger.warning("系统已在运行中")
            return True
        
        # 设置系统
        if not self.setup():
            logger.error("系统设置失败，无法启动")
            return False
        
        # 获取视频流信息
        streams = self.db_manager.get_video_streams()
        if not streams:
            logger.warning("未找到可用的视频流")
            return False
        
        # 启动WebRTC服务器
        if self.webrtc_server:
            logger.info("准备启动WebRTC服务器")
            # 使用线程启动WebRTC服务器
            self.webrtc_thread = threading.Thread(
                target=self._webrtc_server_thread,
                daemon=True
            )
            self.webrtc_thread.start()
            logger.info("WebRTC服务器线程已启动")
        
        # 启动每个视频流处理器
        for stream_info in streams:
            stream_id = str(stream_info['id'])
            
            # 创建并启动流处理器
            processor = StreamProcessor(
                stream_info=stream_info,
                model_config=self.config['model'],
                tracking_config=self.config['tracking'],
                redis_manager=self.redis_manager
            )
            
            if processor.start():
                self.stream_processors[stream_id] = processor
                
                # 如果启用了WebRTC，添加处理后的流
                if self.webrtc_server:
                    name = stream_info.get('name', f"Stream {stream_id}")
                    # 添加处理后的流而不是原始RTSP流
                    self.webrtc_server.add_processed_stream(stream_id, processor, name)
                
                logger.info(f"流 {stream_id} 处理已启动")
            else:
                logger.error(f"流 {stream_id} 处理启动失败")
        
        # 启动监控线程
        self.is_running = True
        self.monitor_thread = threading.Thread(
            target=self._monitor_loop,
            daemon=True
        )
        self.monitor_thread.start()
        
        logger.info(f"系统已启动，处理 {len(self.stream_processors)} 个视频流")
        return True
    
    def stop(self):
        """停止系统"""
        self.is_running = False
        
        # 等待监控线程结束
        if self.monitor_thread and self.monitor_thread.is_alive():
            self.monitor_thread.join(timeout=2.0)
        
        # 停止每个视频流处理器
        for stream_id, processor in self.stream_processors.items():
            processor.stop()
            logger.info(f"流 {stream_id} 处理已停止")
        
        # 清空流处理器
        self.stream_processors.clear()
        
        # 关闭WebRTC服务器
        # 由于WebRTC服务器在单独的线程中运行，它会在程序退出时自动关闭
        # 不需要特别处理
        
        logger.info("系统已停止")
    
    def _monitor_loop(self):
        """监控循环"""
        monitor_interval = self.config['system'].get('monitor_interval', 60)
        
        logger.info(f"监控线程已启动，间隔 {monitor_interval} 秒")
        
        while self.is_running:
            try:
                # 检查数据库中的视频流是否有变化
                self._check_streams()
                
                # 输出每个流的状态
                self._print_status()
                
                # 等待下一次监控
                for _ in range(monitor_interval):
                    if not self.is_running:
                        break
                    time.sleep(1)
                    
            except Exception as e:
                logger.error(f"监控循环出错: {e}")
                time.sleep(10)  # 出错时等待一段时间
    
    def _check_streams(self):
        """检查数据库中的视频流是否有变化"""
        try:
            # 获取最新的视频流信息
            streams = self.db_manager.get_video_streams()
            if not streams:
                return
            
            # 检查新增的流
            current_stream_ids = set(self.stream_processors.keys())
            db_stream_ids = set(str(stream['id']) for stream in streams)
            
            # 需要添加的流
            new_stream_ids = db_stream_ids - current_stream_ids
            for stream_info in streams:
                stream_id = str(stream_info['id'])
                if stream_id in new_stream_ids:
                    # 创建并启动新的流处理器
                    processor = StreamProcessor(
                        stream_info=stream_info,
                        model_config=self.config['model'],
                        tracking_config=self.config['tracking'],
                        redis_manager=self.redis_manager
                    )
                    
                    if processor.start():
                        self.stream_processors[stream_id] = processor
                        
                        # 如果启用了WebRTC，添加处理后的流
                        if self.webrtc_server:
                            name = stream_info.get('name', f"Stream {stream_id}")
                            # 添加处理后的流而不是原始RTSP流
                            self.webrtc_server.add_processed_stream(stream_id, processor, name)
                        
                        logger.info(f"新增流 {stream_id} 处理已启动")
            
            # 需要移除的流
            removed_stream_ids = current_stream_ids - db_stream_ids
            for stream_id in removed_stream_ids:
                if stream_id in self.stream_processors:
                    # 停止并移除流处理器
                    self.stream_processors[stream_id].stop()
                    del self.stream_processors[stream_id]
                    
                    # 如果启用了WebRTC，移除流
                    if self.webrtc_server:
                        self.webrtc_server.remove_stream(stream_id)
                    
                    logger.info(f"移除流 {stream_id} 处理已停止")
            
        except Exception as e:
            logger.error(f"检查视频流变化时出错: {e}")
    
    def _print_status(self):
        """输出系统状态"""
        logger.info(f"系统状态: 运行中={self.is_running}, 处理流数量={len(self.stream_processors)}")
        
        for stream_id, processor in self.stream_processors.items():
            status = processor.get_status()
            logger.info(f"流 {stream_id} 状态: 运行中={status['is_running']}, 帧数={status['frame_count']}, FPS={status['fps']:.1f}, 连接={status['connected']}")


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="多视频流处理系统")
    parser.add_argument("--config", default="config.yaml", help="配置文件路径")
    args = parser.parse_args()
    
    # 创建系统实例
    system = MultiStreamProcessor(args.config)
    
    try:
        # 启动系统
        if system.start():
            logger.info("系统已成功启动，按Ctrl+C停止")
            
            # 等待用户中断
            while True:
                time.sleep(1)
                
        else:
            logger.error("系统启动失败")
            
    except KeyboardInterrupt:
        logger.info("接收到用户中断信号")
    except Exception as e:
        logger.error(f"系统运行出错: {e}")
    finally:
        # 停止系统
        system.stop()
        logger.info("系统已退出")


if __name__ == "__main__":
    main()